package com.edu.chat.web.config;

import cn.hutool.core.collection.CollectionUtil;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.response.ListCollectionsResp;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;

/**
 * @program: KQAIA
 * @description:
 * @author: 无恙
 * @create: 2025-03-27 23:08
 **/

@Component
@Slf4j
public class InitMilvusConfig {

	@Autowired
	private MilvusClientV2 milvusClientV2;

	public static final String COLLECTION_ARTICLE_NAME = "kqaia_article";

	public static final String COLLECTION_DOC_NAME = "kqaia_doc";

	public static final int VECTOR_DIM = 512;

	public static final int topK = 10;

	@PostConstruct
	private void initCollection() {
		//获取所有集合列表
		ListCollectionsResp resp = milvusClientV2.listCollections();

		if(Objects.isNull(resp)){
			log.error("MilvusClientImpl initCollection fail");
			return;
		}

		List<String> collectionNames = resp.getCollectionNames();

		if (CollectionUtil.isEmpty(collectionNames)){
			collectionNames= Collections.emptyList();
		}

		if (!collectionNames.contains(COLLECTION_ARTICLE_NAME)) {
			createArticleCollection();
		}

		if (!collectionNames.contains(COLLECTION_DOC_NAME)) {
			createDocCollection();
		}

	}

	private void createArticleCollection() {

		CreateCollectionReq.CollectionSchema schema = milvusClientV2.createSchema();

		schema.addField(AddFieldReq.builder()
				.fieldName("id")
				.dataType(DataType.Int64)
				.isPrimaryKey(true)
				.autoID(true)
				.build());

		schema.addField(AddFieldReq.builder()
				.fieldName("vector")
				.dataType(DataType.FloatVector)
				.dimension(VECTOR_DIM)
				.build());

		schema.addField(AddFieldReq.builder()
				.fieldName("articleId")
				.dataType(DataType.Int64)
				.build());

		//title
		schema.addField(AddFieldReq.builder()
				.fieldName("title")
				.dataType(DataType.VarChar)
				.maxLength(10000)
				.build());

		//文章内容
		schema.addField(AddFieldReq.builder()
				.fieldName("content")
				.dataType(DataType.VarChar)
				.maxLength(10000)
				.build());

		//知识库ID
		schema.addField(AddFieldReq.builder()
				.fieldName("libraryId")
				.dataType(DataType.Int64)
				.build());

		//租户Id
		schema.addField(AddFieldReq.builder()
				.fieldName("tenantId")
				.dataType(DataType.VarChar)
				.maxLength(32)
				.build());

		// 3.3 Prepare index parameters
		IndexParam indexParamForVectorField = IndexParam.builder()
				.fieldName("vector")
				.indexType(IndexParam.IndexType.IVF_FLAT)
				.metricType(IndexParam.MetricType.COSINE)
				.extraParams(Collections.singletonMap("nlist", 128))
				.build();

		IndexParam indexParamForIdField = IndexParam.builder()
				.fieldName("title")
				.indexType(IndexParam.IndexType.TRIE)
				.metricType(IndexParam.MetricType.COSINE)
				.build();

		List<IndexParam> indexParams = new ArrayList<>();
		indexParams.add(indexParamForIdField);
		indexParams.add(indexParamForVectorField);

		CreateCollectionReq createCollectionReq = CreateCollectionReq.builder()
				.collectionName(COLLECTION_ARTICLE_NAME)
				.collectionSchema(schema)
				.indexParams(indexParams)
				.build();
		milvusClientV2.createCollection(createCollectionReq);

	}

	private void createDocCollection() {
		CreateCollectionReq.CollectionSchema schema = milvusClientV2.createSchema();

		schema.addField(AddFieldReq.builder()
				.fieldName("id")
				.dataType(DataType.Int64)
				.isPrimaryKey(true)
				.autoID(true)
				.build());

		schema.addField(AddFieldReq.builder()
				.fieldName("vector")
				.dataType(DataType.FloatVector)
				.dimension(VECTOR_DIM)
				.build());

		schema.addField(AddFieldReq.builder()
				.fieldName("documentId")
				.dataType(DataType.Int64)
				.build());

		//知识库ID
		schema.addField(AddFieldReq.builder()
				.fieldName("libraryId")
				.dataType(DataType.Int64)
				.build());

		//title
		schema.addField(AddFieldReq.builder()
				.fieldName("text")
				.dataType(DataType.VarChar)
				.maxLength(10000)
				.build());

		//租户Id
		schema.addField(AddFieldReq.builder()
				.fieldName("tenantId")
				.dataType(DataType.VarChar)
				.maxLength(32)
				.build());

		// 3.3 Prepare index parameters
		IndexParam indexParamForVectorField = IndexParam.builder()
				.fieldName("vector")
				.indexType(IndexParam.IndexType.IVF_FLAT)
				.metricType(IndexParam.MetricType.COSINE)
				.extraParams(Collections.singletonMap("nlist", 128))
				.build();

		IndexParam indexParamForIdField = IndexParam.builder()
				.fieldName("documentId")
				.indexType(IndexParam.IndexType.TRIE)
				.metricType(IndexParam.MetricType.COSINE)
				.build();

		List<IndexParam> indexParams = new ArrayList<>();
		indexParams.add(indexParamForIdField);
		indexParams.add(indexParamForVectorField);

		CreateCollectionReq createCollectionReq = CreateCollectionReq.builder()
				.collectionName(COLLECTION_DOC_NAME)
				.collectionSchema(schema)
				.indexParams(indexParams)
				.build();

		milvusClientV2.createCollection(createCollectionReq);
	}

}
